import os
import json
from pycocotools.coco import COCO

def coco_to_yolo_segment(coco_annotation_file, output_dir, image_dir):
    """
    将COCO格式的标注转换为YOLO Segment格式，并生成类别名称文件。
    :param coco_annotation_file: COCO格式的标注文件路径
    :param output_dir: YOLO格式标注文件的输出目录
    :param image_dir: 图像文件目录
    """
    # 加载COCO标注文件
    coco = COCO(coco_annotation_file)
    
    # 获取所有类别信息
    categories = coco.loadCats(coco.getCatIds())
    categories.sort(key=lambda x: x['id'])  # 按类别ID排序
    
    # 生成类别名称文件 classes.txt
    classes_file = os.path.join(output_dir, 'classes.txt')
    with open(classes_file, 'w') as f:
        for cat in categories:
            f.write(f"{cat['name']}\n")
    print(f"Generated {classes_file}")
    
    # 获取所有图像的ID
    image_ids = coco.getImgIds()
    
    # 遍历每张图像
    for img_id in image_ids:
        # 获取图像信息
        img_info = coco.loadImgs(img_id)[0]
        image_width = img_info['width']
        image_height = img_info['height']
        image_filename = img_info['file_name']
        
        # 获取该图像的所有标注
        ann_ids = coco.getAnnIds(imgIds=img_id)
        anns = coco.loadAnns(ann_ids)
        
        # 创建YOLO Segment格式的标注文件
        yolo_annotation_file = os.path.join(output_dir, os.path.splitext(image_filename)[0] + '.txt')
        yolo_images_path = os.path.join(output_dir,"images")
        with open(yolo_annotation_file, 'w') as f:
            for ann in anns:
                # 获取类别ID
                class_id = ann['category_id']
                
                # 获取分割信息（多边形点集）
                segmentation = ann['segmentation']
                
                # COCO的segmentation字段可能包含多个多边形（对于复杂物体）
                # 这里我们只取第一个多边形（如果存在）
                if len(segmentation) == 0 or isinstance(segmentation, list) == False:
                    continue
                
                polygon = segmentation[0]  # 取第一个多边形
                
                # 将多边形点集归一化
                normalized_polygon = []
                for i in range(0, len(polygon), 2):
                    x = polygon[i] / image_width
                    y = polygon[i + 1] / image_height
                    normalized_polygon.extend([x, y])
                
                # 写入YOLO Segment格式的标注
                line = f"{class_id} " + " ".join(map(str, normalized_polygon)) + "\n"
                f.write(line)
    
        print(f"Processed {image_filename}")

# 示例用法
coco_annotation_file = './data/coco/instances_val2017.json'  # COCO格式的标注文件路径
output_dir = './test'  # YOLO Segment格式的标注文件输出目录
image_dir = 'path/to/images'  # 图像文件目录

# 确保输出目录存在
os.makedirs(output_dir, exist_ok=True)

# 转换COCO格式为YOLO Segment格式
coco_to_yolo_segment(coco_annotation_file, output_dir, image_dir)